import spacy
import re
from typing import List

import chromadb
from sentence_transformers import SentenceTransformer

# 加载 spaCy 中文模型（如处理中文可考虑替换为其他分句方式）
nlp = spacy.load("zh_core_web_sm")

def semantic_chunking_with_sliding_window(
    text: str,
    chunk_size: int = 100,
    overlap: int = 10
) -> List[str]:
    """
    基于句子边界进行语义分块 + 滑动窗口策略（使用句子作为overlap单位）
    
    参数:
        text (str): 输入文本
        chunk_size (int): 每个chunk的最大字符数
        overlap (int): 相邻chunk之间的字符级重叠区域（通过句子近似控制）
    
    返回:
        List[str]: 按滑动窗口切分后的文本块列表
    """
    # Step 1: 文本预处理
    text = re.sub(r'\s+', ' ', text).strip()

    # Step 2: 使用 spaCy 分句
    doc = nlp(text)
    sentences = [sent.text.strip() for sent in doc.sents if sent.text.strip()]
    print(len(sentences))

    # Step 3: 滑动窗口切分（按句子长度堆叠）
    chunks = []
    i = 0
    while i < len(sentences):
        current_chunk = []
        current_length = 0
        start_i = i
        while i < len(sentences) and current_length + len(sentences[i]) <= chunk_size:
            current_chunk.append(sentences[i])
            current_length += len(sentences[i])
            i += 1
        chunks.append(' '.join(current_chunk))

        # 回退以实现 overlap（按句子单位回退）
        overlap_len = 0
        j = len(current_chunk) - 1
        while j >= 0 and overlap_len + len(current_chunk[j]) <= overlap:
            overlap_len += len(current_chunk[j])
            j -= 1
        # 回退 i 到重叠起始位置
        i = start_i + j + 1

    return chunks


# 示例用法

sample_text = """自然语言处理（Natural Language Processing, NLP）是计算机科学和人工智能领域的一个重要研究方向，旨在让计算机能够理解、分析、生成和处理人类语言。NLP结合了语言学、计算机科学和数学等多学科的知识，通过算法和模型，使计算机能够从自然语言文本中提取有用的信息，进行语义理解，并执行相关任务。其应用涵盖了从文本分类、情感分析、机器翻译到自动摘要、信息抽取、对话系统等多个领域。
"""

chunks = semantic_chunking_with_sliding_window(sample_text)

for i, chunk in enumerate(chunks):
    print(f"Chunk {i+1} ({len(chunk)} chars):\n{chunk}\n{'-'*60}")


# 1. 初始化向量数据库
client = chromadb.Client()
collection = client.get_or_create_collection("rag_docs")

# 2. 生成嵌入向量
embedder = SentenceTransformer('all-MiniLM-L6-v2')

# 3. 存储分块文本
chunks = semantic_chunking_with_sliding_window(sample_text)
embeddings = embedder.encode(chunks).tolist()

collection.add(
    embeddings=embeddings,
    documents=chunks,
    ids=[f"id{i}" for i in range(len(chunks))]
)

# 4. 检索示例
query = "语言学"
query_embedding = embedder.encode([query]).tolist()

results = collection.query(
    query_embeddings=query_embedding,
    n_results=2
)
print("Top 2 relevant chunks:", results['documents'][0])

# 4. 检索示例
query = "自然语言处理"
query_embedding = embedder.encode([query]).tolist()

results = collection.query(
    query_embeddings=query_embedding,
    n_results=2
)
print("Top 2 relevant chunks:", results['documents'][0])


# 查询向量数据库中数据条数
num_docs = collection.count()
print(f"向量数据库中共有 {num_docs} 条数据。")

# 查询全部数据
all_data = collection.get(include=["documents", "embeddings"])

documents = all_data["documents"]
embeddings = all_data["embeddings"]
ids = all_data["ids"]

print("\n全部数据（仅显示前几项预览）:")
for i in range(min(5, len(documents))):  # 控制输出数量避免过长
    print(f"ID: {ids[i]}")
    print(f"文档内容: {documents[i]}")
    print(f"嵌入向量（前5维）: {embeddings[:5]}")   
    print("-" * 60)

# 嵌入维度（假设所有向量维度相同）
try:
    embedding_dims = len(embeddings[0])
    print(f"嵌入向量维度：{embedding_dims}")
except Exception as e:
    print(f"无法计算嵌入向量维度，错误信息: {e}")

